import numpy as np
import matplotlib.pyplot as plt
import math

# 定义一个函数，用于生成正弦波
def getSin(amp, freq, phase, sampleList):
    return amp * np.sin(-2 * math.pi * freq * sampleList + phase)

# 定义一个函数，用于生成余弦波
def getCos(amp, freq, phase, sampleList):
    return amp * np.cos(-2 * math.pi * freq * sampleList + phase) 

# 定义去噪函数，将振幅小于阈值的噪声去除
def denoise(arr, thresh):
    mask = arr > thresh
    return mask * arr

# 1. 获得混合波形
srate=3000
t=np.linspace(0,1,srate)
s1 = getSin(amp=1.5, freq=30, phase=0, sampleList=t)
s2 = getCos(amp=3, freq=5, phase=0, sampleList=t)
s3 = getSin(amp=10, freq=100, phase=0, sampleList=t)
s4 = getCos(amp=20, freq=120, phase=0, sampleList=t)
# 将四个正弦和余弦波相加，得到混合信号
m = s1 + s2 + s3 + s4
# 2. 获得傅里叶系数
fCoefs = np.fft.fft(m,srate)

# 3. 获得振幅列表：每一个绕线的重心到原点的距离
amp_list=2 * np.abs(fCoefs / srate)

# 把频率轴从0~1000 转变成 0~499 然后 -500~-1
freqs = np.fft.fftfreq(len(amp_list), 1/srate)
# 然后把 频率轴 和 数据 都变成 0hz 在中间，向左是负频率，向右是正频率的形式
amp_shifted=np.fft.fftshift(amp_list)
freq_shift=np.fft.fftshift(freqs)
# 创建两个子图，分别用于显示原始振幅频谱和处理后的振幅频谱
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
# 绘制原始振幅频谱
ax[0].stem(freq_shift, amp_shifted)
ax[0].set_xlim([-150, 150])
#ax[0].set_ylim([-1, 21])
ax[0].grid()
# 对振幅频谱进行去噪处理
# 1. 去掉所有低振幅噪音
# mask = np.ones_like(amp_shifted) #创建一个全1数组作为掩膜
# mask[np.abs(freq_shift) < 1] = 0 #将频率小于1的掩膜位置设为0
# amp_shifted *= mask   #去除频率小于1的所有噪音
amp_shifted = denoise(amp_shifted, 1)
# 2. 去掉所有的高于110hz的正负频率噪音
amp_shifted[(freq_shift > 110) | (freq_shift < -110)] = 0
# 绘制处理后的振幅频谱
ax[1].stem(freq_shift, amp_shifted)
# 设置第二个子图的坐标轴范围
ax[1].set_xlim([-150, 150])
ax[1].set_ylim([-1, 21])
# 给第二个子图添加网格线
ax[1].grid()
# 显示图形
plt.show()